﻿# backtest/engine.py
import pandas as pd
import akshare as ak

class AnimatedBacktestEngine:
    def __init__(self):
        self.data = None
        self.current_idx = 10
        self.strategy = None
        self.position = 0
        self.cash = 100000
        self.holdings = 0
        self.trades = []

    def add_data(self, symbol: str, start_date: str, end_date: str):
        try:
            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date,
                end_date=end_date,
                adjust="qfq"
            )
            if df is None or df.empty:
                return pd.DataFrame()

            df.rename(columns={
                '日期': 'datetime',
                '开盘': 'open',
                '最高': 'high',
                '最低': 'low',
                '收盘': 'close',
                '成交量': 'volume'
            }, inplace=True)
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)

            df['ma5'] = df['close'].rolling(5).mean()
            df['ma10'] = df['close'].rolling(10).mean()

            self.data = df.dropna()
            self.current_idx = 10
            return self.data
        except Exception as e:
            print(f"数据获取失败: {e}")
            return pd.DataFrame()

    def add_strategy(self, strategy):
        self.strategy = strategy

    def run_next(self):
        if self.current_idx >= len(self.data) - 1:
            return False

        self.current_idx += 1
        current_bar = self.data.iloc[self.current_idx]
        price = current_bar['close']

        if self.strategy:
            signal = self.strategy.generate_signal(self.data, self.current_idx)

            if signal == 'buy' and self.cash > price:
                shares = int(self.cash / price)
                self.holdings += shares
                self.cash -= shares * price
                self.trades.append({
                    'type': 'buy',
                    'price': price,
                    'shares': shares,
                    'time': current_bar.name
                })

            elif signal == 'sell' and self.holdings > 0:
                self.cash += self.holdings * price
                self.trades.append({
                    'type': 'sell',
                    'price': price,
                    'shares': self.holdings,
                    'time': current_bar.name
                })
                self.holdings = 0

        return True

    def get_current_data(self):
        return self.data.iloc[:self.current_idx+1]

    def get_trades(self):
        return pd.DataFrame(self.trades) if self.trades else pd.DataFrame(columns=['type', 'price', 'shares', 'time'])

    def get_stats(self):
        if self.current_idx >= len(self.data):
            return {}
        current_price = self.data.iloc[self.current_idx]['close']
        value = self.cash + self.holdings * current_price
        return {'value': value, 'cash': self.cash, 'holdings': self.holdings}